🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20686

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit a551ca3 with merge base f41d66d ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Hi @psiddh,
This PR is to optimize performance for PCQ embedding on HTP.
Could you please have a look?
Thanks,
Hutton

@claude Can you review this PR ?

Claude finished @psiddh's task in 3m 23s —— View job

PR Review: Optimize performance of PCQ embedding

• Read the PR diff and changed files
• Analyze op_embedding.py restructure
• Analyze test and llm_wrappers.py changes
• Post review

Nice cleanup — the new pattern is clearly better than the old one. Here's my review.

Summary of the change

The core change in op_embedding.py inverts the op order for per-channel-quantized (PCQ) embeddings:

• Before: Gather(per-channel weight) → intermediate output where the gather output was forced to a per-tensor approximation (scale = max(scales), zp = max(zero_points)), then a Convert requantized to the activation encoding. That max-scale approximation is exactly what caused the accuracy problem the test was skipped for ("Bad accuracy for HTP").
• After: Convert(per-channel weight) → per-tensor activation-encoded weight table → Gather → output. The per-channel→per-tensor requant now happens on the weight table itself (properly, per-channel), and gather just indexes it. This removes the lossy max(scales) heuristic entirely, which is why the test can now be enabled.

This is a solid structural improvement, and dropping the node.meta[QCOM_QUANT_ATTRS] mutate-and-restore dance (node_quant_attrs/intermediate_quant_attrs) removes a fragile side effect. Removing the now-unused QCOM_QUANT_MAX/MIN/SCALE(S)/ZERO_POINT(S) imports is good hygiene.

Correctness notes (looks right)

• The final output wrapper is now registered under node.name (op_embedding.py:96-103), so downstream consumers resolve the correct tensor. The Convert output is an intermediate _convert native tensor feeding into Gather. ✅
• Gather input (converted weight) and Gather output share the same activation encoding via get_quant_encoding_conf(node, node) — correct, since gather only indexes and must not change encoding. ✅
• op_wrapper_list ordering is correct: convert_op is appended before gather_op, so the convert executes first. ✅

Questions / things worth confirming

1. Memory/perf of converting the full weight table (op_embedding.py:71-89)
The Convert now runs over the entire embedding table (dims=weight_tensor.size(), i.e. vocab × hidden) instead of over the small gather output. For LLM-sized vocabularies this is a large tensor. Since the convert input is a QNN_TENSOR_TYPE_STATIC weight, I assume the intent is that HTP constant-folds this convert at compile time so only the gather runs at inference — which would make this a genuine perf win. Can you confirm the convert-of-static is folded (and doesn't materialize a full-size runtime NATIVE tensor each inference)? A one-line note in the code or PR would help future readers understand why this is faster.

2. llm_wrappers.py getattr/setattr asymmetry (llm_wrappers.py:785-797)
The new _get_attr correctly traverses dotted targets (e.g. static_llama.tok_embedding.weight) for the read side. But the write side still uses a flat setattr(unquantized_model, unquantized_node.target, ...). If unquantized_node.target can also be dotted, setattr will create a single literal attribute named "a.b.c" on the model rather than setting model.a.b.c. Is unquantized_node.target guaranteed to be flat here, while only quantized_node.target is nested? If not, the write path needs the same traversal.

3. Version gate
is_qnn_sdk_version_less_than("2.48") only compares within the same major version (current_major == target_major and current_minor < target_minor). So on a hypothetical QNN 3.x the check returns False and the test runs — which is the desired behavior here, just flagging that the helper is major-scoped by design.

Minor

• Test now covers both use_16a8w and use_16a4w with subTest — good added coverage. Consider adding a short comment noting the QNN 2.48 requirement mirrors the backend change (the skip message "UT pass after QNN 2.48" covers it, so optional).

Overall this looks like a clean, well-motivated change. My main ask is confirming point #1 (that the static convert is folded so this is actually faster, not just more accurate) and clarifying the setattr asymmetry in #2.

One qq: The test is skipped below QNN 2.48. does that mean the new pcq-embedding pattern only works on 2.48+,
or does it still work on 2.37 and 2.48 is just where the accuracy test passes?

Also just for my understanding : The convert now runs over the whole weight table instead of just the gathered rows. Since the weight is static, is this convert constant-folded during QNN compile (so it's computed once and the runtime
graph is just the gather)? Just want to confirm it's a one-time AOT cost and not a per-inference convert of the full table.

One qq: The test is skipped below QNN 2.48. does that mean the new pcq-embedding pattern only works on 2.48+, or does it still work on 2.37 and 2.48 is just where the accuracy test passes?

Yes, this new optimize pattern is only supported after QNN 2.48. Otherwise, it will failed to compile

Also just for my understanding : The convert now runs over the whole weight table instead of just the gathered rows. Since the weight is static, is this convert constant-folded during QNN compile (so it's computed once and the runtime graph is just the gather)? Just want to confirm it's a one-time AOT cost and not a per-inference convert of the full table.

Yes, it will be fused by QNN HTP during AOT.

One qq: The test is skipped below QNN 2.48. does that mean the new pcq-embedding pattern only works on 2.48+, or does it still work on 2.37 and 2.48 is just where the accuracy test passes?

Yes, this new optimize pattern is only supported after QNN 2.48. Otherwise, it will failed to compile

Thanks for confirming. Since it fails to compile below 2.48 and the builder emits this pattern unconditionally, won't pcq embedding now break for anyone on our pinned 2.37? Should we either gate the builder on is_qnn_sdk_version_less_than("2.48") (fall back to the old pattern below 2.48), or bump the minimum QNN version to 2.48 and move CI accordingly? Otherwise the test skip hides a compile failure for 2.37 users

Thanks for confirming. Since it fails to compile below 2.48 and the builder emits this pattern unconditionally, won't pcq embedding now break for anyone on our pinned 2.37? Should we either gate the builder on is_qnn_sdk_version_less_than("2.48") (fall back to the old pattern below 2.48), or bump the minimum QNN version to 2.48 and move CI accordingly? Otherwise the test skip hides a compile failure for 2.37 users

Thanks for flagging that. I agree — I keep the legacy pattern for pcq embedding to maintain backward compatibility.